#include"procmonitor.h"

extern "C" {
  #include"options.h"
  #include"proc.h"
}


#include<stdlib.h>
#include<pwd.h>
#include<dirent.h>
#include<errno.h>
#include<assert.h>

#include<vector>
#include<list>
#include<string>
#include<map>
#include<algorithm>

using namespace std;

char * _linebuf;
size_t _linebufsize;

/* kernel version major, minor, patch*/
int kv_maj;
int kv_min;
int kv_pat;

/* the age of this process */
int age;

int cpu_number;

map<uid_t, string> user_map;
map<uid_t, map<pid_t, ProcEnt *> > user_proc_map;

static void record_user( uid_t uid )
{
  struct passwd * pwd = getpwuid( uid );
  if( pwd == NULL )
    {
      fprintf( stderr, "Failed to lookup user with id \"%d\", aborted.\n", uid );
      exit( 1 );
    }
  user_map.insert( map<uid_t,string>::value_type( uid, string( pwd->pw_name ) ) );
}

static void record_user( char * uname )
{
  // lookup uid
  struct passwd * pwd = getpwnam( uname );
  if( pwd == NULL )
    {
      fprintf( stderr, "Failed to lookup user with name \"%s\", aborted.\n", uname );
      exit( 2 );
    }
  user_map.insert( map<uid_t,string>::value_type( pwd->pw_uid, string( uname ) ) );
}

static void lookup_uid()
{
  if( options.un != NULL )
    {
      if( options.un[ 0 ] != '@' )
        {
          record_user( options.un );
        }
      else
        {
          FILE * fp = fopen( options.un + 1, "r" );
          if( fp == NULL )
            {
              fprintf( stderr, "Failed to file open \"%s\": [%d]: %s\n", options.un + 1, errno, strerror( errno ) );
              exit( 3 );
            }
          char namebuf[ 64 ];
          while( fscanf( fp, "%60s", namebuf ) == 1 )
            {
              record_user( namebuf );
            }
          fclose( fp );
        }
    }
  else if( options.uid != NULL )
    {
      if( options.uid[ 0 ] != '@' )
        {
          char * q;
          uid_t u = strtol( options.uid, &q, 10 );
          if( *q != '\0' )
            {
              fprintf( stderr, "Illegal value of --user-id, aborted.\n" );
              exit( 4 );
            }
          record_user( u );
        }
      else
        {
          FILE * fp = fopen( options.uid + 1, "r" );
          if( fp == NULL )
            {
              fprintf( stderr, "Failed to file open \"%s\": [%d]: %s\n",
                       options.uid + 1, errno, strerror( errno ) );
              exit( 5 );
            }
          uid_t id;
          while( fscanf( fp, "%d", &id ) == 1 )
            {
              record_user( id );
            }
          fclose( fp );
        }
    }
  else
    {
      record_user( getuid() );
    }
}

class ProcFilter
{
public:
  ProcFilter( uid_t _uid ): uid( _uid ) {}
  ~ProcFilter() {} 
  uid_t uid;
  bool operator ()(uid_t _uid)
  {
    return uid == _uid;
    // return true;
  }
};

template<class ProcFilter>
static ProcEnt * gen_proc_ent( map<pid_t, ProcEnt *> & proc_map, pid_t pid, ProcFilter filter )
{
  char numberbuf[ 16 ];
  if( snprintf( numberbuf, 16, "%d", pid ) >= 16 )
    {
      proc_map.erase( pid );
      return NULL;
    }
  string proc_dir = string( "/proc/" ) + string( numberbuf );
  string proc_status = proc_dir + string( "/status" );
  FILE * fp = fopen( proc_status.c_str(), "r" );
  if( fp == NULL )
    {
      proc_map.erase( pid );
      return NULL;
    }
  int r;
  uid_t ruid, euid, suid, fuid;
  while( ( r = getline( &_linebuf, &_linebufsize, fp ) ) != -1 )
    {
      if( sscanf( _linebuf, "Uid: %d %d %d %d", &ruid, &euid, &suid, &fuid ) == 4 )
        {
          break;
        }
    }
  fclose( fp );

  if( !filter( ruid ) && !filter( euid ) && !filter( suid ) && !filter( fuid ) )
    {
      proc_map.erase( pid );
      return NULL;
    }

  string cmdline_file( proc_dir + "/cmdline" );
  if( ( fp = fopen( cmdline_file.c_str(), "r" ) ) == NULL )
    {
      proc_map.erase( pid );
      return NULL;
    }

  ProcEnt * & proc = proc_map[ pid ];
  if( proc == NULL )
    {
      proc = new ProcEnt( pid, proc_dir, ruid, euid, suid, fuid );
    }

  vector<char> word(128);
  word.clear();
  int c;
  while( ( c = fgetc( fp ) ) != EOF )
    {
      if( word.size() == 0 && c == 0 )
        {
          break;
        }
      word.push_back( (char)c );
      if( c == 0 )
        {
          string s = string( (char *)&word[ 0 ] );
          proc->cmd_vec.push_back( s );
          word.clear();
        }
    }
  fclose( fp );

  if( proc->cmd_vec.size() == 0 )
    {
      // fprintf( stderr, "warning: proc cmd array is null\n" );
      proc_map.erase( pid );
      delete proc;
      proc = NULL;
    }
  else
    {
      proc->age = age;
    }

  return proc;
  
}

static int read_statm( ProcEnt * proc )
{
  FILE * fp;
  if( ( fp = fopen( proc->proc_statm.c_str(), "r" ) ) == NULL )
    {
      return 1;
    }
  
  if( fscanf( fp, "%d %d %d", &proc->mem_vir, &proc->mem_res, &proc->mem_shr ) !=3 )
    {
      proc->mem_vir = 0;
      proc->mem_res = 0;
      proc->mem_shr = 0;
      fclose( fp );
      return 1;
    }
  proc->mem_vir <<= 2;
  proc->mem_res <<= 2;
  proc->mem_shr <<= 2;
  fclose( fp );
  return 0;
}

static int read_jiffies( ProcEnt * proc )
{
  FILE * fp;
  if( ( fp = fopen( proc->proc_stat.c_str(), "r" ) ) == NULL )
    {
      proc->record_jiffy( 0, 0 );
      return 1;
    }

  if( getline( &_linebuf, &_linebufsize, fp ) != -1 )
    {
      proc->record_jiffy_time();
      char * p = _linebuf;
      char * q = p;
      int item = 0;
      long int k[ 2 ];
      while( p && *p )
        {
          if( *p == ' ' || *p == '\t' )
            {
              ++item;
              if( item == 14 || item == 15 )
                {
                  *p = '\0';
                  char * tmp;
                  k[ item - 14 ] = strtol( q, &tmp, 10 );
                  if( tmp != p )
                    {
                      k[ item - 14 ] = 0;
                    }
                  if( item == 15 )
                    {
                      proc->record_jiffy( k[ 0 ], k[ 1 ] );
                      break;
                    }
                }
              ++p;
              q = p;
            }
          else
            {
              ++p;
            }
        }
    }
  fclose( fp );
  long int ju,jk;
  proc->get_jiffy( &ju, &jk );
  // fprintf( stdout, "%d %d %s: %lu %lu\n", proc->born, proc->age, proc->cmd_vec[ 0 ].c_str(), ju, jk );
  return 0;
}

int gather_proc_information_for_user( uid_t uid, map<pid_t, ProcEnt *> & proc_map )
{
  DIR * d;
  if( ( d = opendir( "/proc/" ) ) == NULL )
    {
      fprintf( stderr, "opendir failed: [%d]: %s\n", errno, strerror( errno ) );
      fflush( stderr );
      exit( 6 );
    }
  struct dirent * de;
  while( ( de = readdir( d ) ) != NULL )
    {
      char * p;
      pid_t pid = strtol( de->d_name, &p, 10 );
      if( *p == 0 )
        {
          ProcEnt * procent = gen_proc_ent( proc_map, pid, ProcFilter(uid) );
          if( procent )
            {
              read_statm( procent );
              read_jiffies( procent );
              proc_map.insert( map<pid_t, ProcEnt *>::value_type( procent->pid, procent ) );
            }
        }
    }
  closedir( d );

  map<pid_t, ProcEnt *> tmp_map;
  tmp_map.swap( proc_map );
  map<pid_t, ProcEnt *>::iterator i = tmp_map.begin();
  while( i != tmp_map.end() )
    {
      if( i->second->age >= ::age )
        {
          proc_map.insert( map<pid_t, ProcEnt *>::value_type( i->first, i->second ) );
        }
      ++i;
    }
  assert( proc_map.size() <= tmp_map.size() );
}

void check_kernel_version()
{
  int err;
  if( get_kernel_version( &kv_maj, &kv_min, &kv_pat, &err ) == 0 )
    {
      if( kv_maj != 2 || ( kv_maj == 2 && kv_min < 6 ) )
        {
          fprintf( stderr, "Unsupported kernel version: %d.%d.%d\n", kv_maj, kv_min, kv_pat );
          fprintf( stderr, "Kernel version should >= 2.6\n" );
          exit( 7 );
        }
    }
  else
    {
      fprintf( stderr, "failed to get kernel version: [%d]: %s\n", err, strerror( err ) );
      exit( 8 );
    }
}

class MyLarge
{
public:
  bool operator() ( ProcEnt * p1, ProcEnt * p2 )
  {
    return p1->_last_cpu_percentage > p2->_last_cpu_percentage;
  }
};

int report_proc_for_user( uid_t uid, map<pid_t, ProcEnt *> & proc_map )
{
  int vir = 0, res = 0, shr = 0;
  map<pid_t, ProcEnt *>::iterator it = proc_map.begin(), end = proc_map.end();
  double per = 0;
  vector<ProcEnt *> _tmp_vec;
  while( it != end )
    {
      vir += it->second->mem_vir;
      res += it->second->mem_res;
      shr += it->second->mem_shr;
      if( it->second->age - it->second->born >= 1 )
        {
          long int ju, jk;
          it->second->get_jiffy( &ju, &jk );
	  int duration = it->second->get_duration();
	  // per += (double)( ju + jk ) * 10 * 10000 / ( duration * cpu_number );
          double _p = (double)( ju + jk ) * 10000 / ( (double)duration * HZ / 1000 * cpu_number );
          if( options.report_cpu_hogs )
            {
              it->second->_last_cpu_percentage = _p / 100;
              _tmp_vec.push_back( it->second );
            }
          per += _p;
        }
      ++it;
    }
  if( per >= 1 )
    {
      per /= 100;
      time_t now;
      time( &now );
      struct tm nowtm;
      localtime_r( &now, &nowtm );
      fprintf( stdout, "%4d%02d%02d%02d%02d%02d %s[%d] %dk(%dM) %dk(%dM) %dk(%dM) %.2f%%\n",
               nowtm.tm_year + 1900, nowtm.tm_mon+1, nowtm.tm_mday,
               nowtm.tm_hour, nowtm.tm_min, nowtm.tm_sec, 
               user_map[ uid ].c_str(), uid, 
               vir, vir/1024, res, res/1024, shr, shr/1024,
               per );
      if( options.report_cpu_hogs && _tmp_vec.size() > 0 )
        {
          vector<ProcEnt *>::iterator b = _tmp_vec.begin(), e = _tmp_vec.end();
          sort( b, e, MyLarge() );
          for( b = _tmp_vec.begin(); b != _tmp_vec.end(); ++b )
            {
              ProcEnt * proc = *b;
              if( proc->_last_cpu_percentage / per * 100 >= options.cpu_hog_threshold )
                {
                  fprintf( stdout, "    CPU(%.2f%% %.2f%%) %s\n",
                           proc->_last_cpu_percentage, proc->_last_cpu_percentage / per  * 100,
                           proc->cmd_vec[ 0 ].c_str() );
                }
              else
                {
                  break;
                }
            }
          _tmp_vec.clear();
          // map<pid_t, ProcEnt *>::iterator it = proc_map.begin(), end = proc_map.end();
          // while( it != end )
          //   {
          //     ProcEnt * proc = it->second;
          //     if( per >= 0.1 )
          //       {
          //         if( proc->_last_cpu_percentage / per * 100 >= options.cpu_hog_threshold )
          //           {
          //             fprintf( stdout, "    CPU(%.2f%% %.2f%%) %s\n",
          //                      proc->_last_cpu_percentage, proc->_last_cpu_percentage / per  * 100,
          //                      proc->cmd_vec[ 0 ].c_str() );
          //           }
          //       }
          //     ++it;
          //   }
        }
    }
  else
    {
      // percentage <= 1%, just ignore
    }
  fflush( stdout );
}

int main( int argc, char ** argv )
{
  _linebuf = 0;
  _linebufsize = 0;
  if( process_options( argc, argv ) )
    {
      exit( 9 );
    }
  int err;
  check_kernel_version();
  cpu_number = get_cpu_number( &err );
  
  fprintf( stdout, "Linux: %d.%d.%d\n", kv_maj, kv_min, kv_pat );
  fprintf( stdout, "CPU(s): %d\n", cpu_number );
  fprintf( stdout, "HZ: %d\n", HZ );
  fflush( stdout );
  if( cpu_number <= 0 )
    {
      return 1;
    }

  
  lookup_uid();
  map<uid_t, string>::iterator i = user_map.begin(), e = user_map.end();
  fprintf( stdout, "User:" );
  for( ; i != e; ++i )
    {
      fprintf( stdout, " %s[%d]", i->second.c_str(), i->first  );
      user_proc_map[ i->first ] = map<pid_t, ProcEnt *>();
    }
  fprintf( stdout, "\n" );
  fflush( stdout );
  
  age = 0;
  while( 1 )
    {
      ++age;
      map<uid_t, map<pid_t, ProcEnt *> >::iterator
        i = user_proc_map.begin(), e= user_proc_map.end();
      for( ; i != e; ++i )
        {
          gather_proc_information_for_user( i->first, i->second );
          report_proc_for_user( i->first, i->second );
        }
      sleep( options.interval );
    }
  return 0;
}


ProcEnt::ProcEnt( pid_t _pid, string & _proc_dir, pid_t _ruid, pid_t _euid, pid_t _suid, pid_t _fuid )
  : pid( _pid ), proc_dir( _proc_dir ), proc_statm( _proc_dir + "/statm" ), proc_stat( _proc_dir + "/stat" ),
    ruid( _ruid ), euid( _euid ), suid( _suid ), fuid( _fuid ), mem_vir(0), mem_res(0), mem_shr(0),
    current_jiffy( 0 ), born(::age), age( ::age )
{

}

ProcEnt::~ProcEnt() {}
 
